// modelUtils.js
const fs = require('fs');
const path = require('path');
const { generateShortUUID } = require('./uuidUtils');
const { DatasetListDB } = require('../dbs/dbDatasetList');
const BaseDatasetWrapper = require('../labels/BaseDatasetWrapper');

class ModelUtils {
    constructor(app) {
        this.app = app;
        this.current_train_id = null;
    }

    // 初始化训练所需要的基本信息
    async init(train_id) {
        try {
            // 获取对应的训练信息，包括model_id和dataset_id
            this.train_id = train_id;

            const local_models = await this.app.db.models_local.GetLocalModels(this.train_id);
            if (local_models.length === 0) {
                return false;
            }
            this.train_data = local_models[0];
            this.model_id = this.train_data.model_id;
            this.dataset_id = this.train_data.dataset_id;
            // console.log(this.train_data);

            // 获取模型信息，包括名称、训练脚本、评测脚本等
            const cloud_models = await this.app.models_cloud_db.models_cloud.GetCloudModels(this.model_id);
            if (cloud_models.length === 0) {
                return false;
            }
            this.cloud_model = cloud_models[0];

            // 获取数据集信息，包括名称、数据集路径等
            this.dataset = await this.app.db.datasets.getDatasetInfo(this.dataset_id);

            return true;
        } catch (err) {
            console.error(`Error get local models: `, err);
            return false;
        }
    }

    // 生成训练所需要的配置文件
    async generateParamsConfig(dataset_id, output_dir, train_cfg) {
        return new Promise(async (resolve, reject) => {
            try {
                // 获取dataset_id对应的数据集信息
                let dataset_info = await this.app.db.datasets.getDatasetInfo(dataset_id);
                // console.log(dataset_info);

                // 打开详情数据库，获取对应的subsets信息
                let dblist = this.app.list_dbs[dataset_id];
                if (!dblist) {
                    this.app.list_dbs[dataset_id] = new DatasetListDB(this.app, dataset_id);
                    dblist = this.app.list_dbs[dataset_id];
                }

                // 获取subset列表
                let rows = await dblist.subsets.get();
                let subsets_info = rows.filter((subset) => subset.subset_type !== 'directory');
                // console.log(`subsets_info = `, subsets_info);

                // 调用对应数据的导出方法
                let total_config = train_cfg;
                total_config['dataset_config'] = await BaseDatasetWrapper.invokeExport(this.app, train_cfg, dataset_info, subsets_info);

                // 将config内容保存到output目录下
                if (!fs.existsSync(output_dir)) {
                    fs.mkdirSync(output_dir, { recursive: true });
                }
                const outputFilePath = path.join(output_dir, `params.json`);
                fs.writeFileSync(outputFilePath, JSON.stringify(total_config, null, 2));

                // 返回对应的文件路径
                resolve(outputFilePath);
            } catch (error) {
                console.error(`Error generating config: `, error);
                reject(error);
            }
        });
    }

    // 开始训练模型
    // train_id： 对应的本地模型id
    // train_cfg: 训练的配置信息
    async startTrain(train_id, train_cfg) {
        try {
            if (!train_id) {
                return { code: -1, msg: 'train_id不能为空' };
            }

            // 判断当前是否有正在运行的任务
            if (this.app.tasks.processing) {
                return { code: -1, msg: '队列中有任务正在运行' };
            }

            // 获取训练任务的信息
            const initSuccess = await this.init(train_id);
            if (!initSuccess) {
                return { code: -1, msg: '获取训练任务信息失败' };
            }

            let remoteId = this.train_data.remote_id;
            let envName = this.train_data.env_name;
            console.log(`[remoteId = ${remoteId}, envName = ${envName}]`);
            console.log(`ModelTrainer init model_id = ${this.model_id}, dataset_id = ${this.dataset_id}`);
            // console.log(`ModelTrainer init cloud_model = `, this.cloud_model);
            // console.log(`ModelTrainer init dataset = `, this.dataset);

            // 判断训练集是否存在
            if (!this.dataset) {
                return { code: -1, msg: '训练集不存在' };
            }

            // 准备训练脚本的参数
            let script_train = this.cloud_model.script_train;
            let taskId = generateShortUUID();
            let output_dir = path.join(this.app.userDataPath, 'output', taskId);
            let params_path = await this.generateParamsConfig(this.dataset_id, output_dir, train_cfg);
            console.log(`python "${script_train}" "${params_path}" "${output_dir}"`);

            // 将任务添加到队列中，开始正式训练
            this.current_train_id = train_id;
            let record_id = await this.app.db.train_records.insert({
                local_model_id: this.train_id,
                status: 'init',
                task_id: taskId,
                output_dir: output_dir,
                train_cfg: JSON.stringify(train_cfg)
            });

            this.app.tasks.add(taskId, async (setProcess) => {
                // 执行python脚本
                let result = null;
                try {
                    result = await this.app.pythonEnv.runCommandEx(envName, 'python', [`"${script_train}"`, `"${params_path}"`, `"${output_dir}"`], setProcess);
                } catch (error) {
                    this.app.queue.enqueue(taskId, {
                        event: 'stdio',
                        params: {
                            logData: error.msg
                        }
                    });
                    result = { code: -1, msg: '子进程执行异常错误' };
                } finally {
                    // 处理子进程任务状态并异步返回结果
                    // console.log('runCommandEx completed, result = ', result);
                    this.current_train_id = null;
                    let status = result.code === 0 ? 'success' : 'failed';
                    await this.app.db.train_records.update({
                        id: record_id,
                        status: status
                    });

                    // 返回结果，可以考虑通过websocket或者消息队列返回
                    this.app.queue.enqueue(taskId, {
                        event: 'result',
                        params: {
                            code: result.code,
                            msg: record_id
                        }
                    });

                    return result;
                }
            });

            // 开始监控文件的变化，pageId和train_id作为回调的信息
            let jsproxy_path = this.cloud_model.script_jsproxy;
            const handlerClass = require(jsproxy_path);
            const jsproxy = new handlerClass();
            this.app.fileWatcher.addWatcher(output_dir, taskId, jsproxy, train_cfg);

            // 更新训练任务状态
            await this.app.db.train_records.update({
                id: record_id,
                status: 'running'
            });

            return { code: 0, msg: { method: 'async', task_id: taskId } };
        } catch (error) {
            console.log(`startTrain error = `, error);
            return { code: -1, msg: '创建训练任务失败（未知错误）' };
        }
    }

    // 终止某一个本地模型的所有训练任务
    async stopTrain(train_id) {
        console.log(`stopTrain train_id = `, train_id);
        if (!train_id) {
            return { code: -1, msg: 'train_id不能为空' };
        }

        try {
            let records = await this.app.db.train_records.getRecords(train_id);

            for (let i = 0; i < records.length; i++) {
                let record = records[i];

                if (record.status === 'running') {
                    // 判断当前是否有任务正在运行
                    let taskId = record.task_id;
                    let taskStatus = this.app.tasks.getTaskStatus(taskId);
                    if (!taskStatus) {
                        console.log(`task ${taskId} not found`);
                        // return { code: -2, msg: '队列中未找到正在运行的任务' };
                    }

                    // 停止训练任务
                    if ((await this.app.tasks.stop(taskId)) == false) {
                        console.log(`stop task ${taskId} failed`);
                        // return { code: -3, msg: '停止训练任务失败' };
                    }

                    // 更新训练任务状态
                    await this.app.db.train_records.update({
                        id: record.id,
                        status: 'stopped'
                    });
                }
            }

            this.current_train_id = null;
            return { code: 0, msg: '停止训练任务成功' };
        } catch (error) {
            console.log(`stopTrain error = `, error);
            return { code: -1, msg: '停止训练任务失败（未知错误）' };
        }
    }

    // 开始评测模型
    // record_id: 对应的训练记录，方便查找到对应的目录
    // model_path: 模型的路径
    async startEval(record_id, model_path) {
        // 判断当前是否有正在运行的任务
        if (this.app.tasks.processing) {
            return { code: -1, msg: '队列中有任务正在运行' };
        }

        // 获取训练记录以及训练任务的信息
        let records = await this.app.db.train_records.getRecord(record_id);
        if (records.length === 0) {
            return { code: -1, msg: `Record with ID ${record_id} not found` };
        }

        let record = records[0];
        let train_cfg = record.train_cfg;
        let output_dir = record.output_dir;
        let train_id = record.local_model_id;

        // 获取模型信息
        const initSuccess = await this.init(train_id);
        if (!initSuccess) {
            return { code: -1, msg: '获取训练任务信息失败' };
        }
        let remoteId = this.train_data.remote_id;
        let envName = this.train_data.env_name;
        console.log(`[remoteId = ${remoteId}, envName = ${envName}]`);
        console.log(`ModelTrainer init model_id = ${this.model_id}, dataset_id = ${this.dataset_id}`);

        let script_eval = this.cloud_model.script_eval;
        let params_path = path.join(output_dir, 'params.json');
        let eval_out_dir = path.join(output_dir, 'test');
        let taskId = generateShortUUID();
        console.log(`python "${script_eval}" "${params_path}" "${model_path}" "${eval_out_dir}"`);

        // 启动python环境评测模型
        this.app.tasks.add(taskId, async (setProcess) => {
            let result = null;
            try {
                result = await this.app.pythonEnv.runCommandEx(envName, 'python', [`"${script_eval}"`, `"${params_path}"`, `"${model_path}"`, `"${eval_out_dir}"`], setProcess);
            } catch (error) {
                this.app.queue.enqueue(taskId, {
                    event: 'stdio',
                    params: {
                        logData: error.msg
                    }
                });
                result = { code: -1, msg: '子进程执行异常错误' };
            } finally {
                // 返回结果，可以考虑通过websocket或者消息队列返回
                // console.log(`startEval result = `, result);
                this.app.queue.enqueue(taskId, {
                    event: 'result',
                    params: {
                        code: result.code,
                        msg: {
                            taskId,
                            record_id,
                            model_path
                        }
                    }
                });
                return result;
            }
        });

        return { code: 0, msg: { method: 'async', task_id: taskId } };
    }

    // 终止某一个本地模型的所有评测任务
    async stopEval(record_id) {}

    async clear() {
        if (this.current_train_id) {
            let result = await this.stopTrain(this.current_train_id);
            console.log('clear training record, result = ', result);
            this.current_train_id = null;
        }
    }
}

module.exports = { ModelUtils };
